package com.wkh.project.manager;

import com.wkh.project.common.ErrorCode;
import com.wkh.project.exception.BusinessException;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

@Component
public class AIManager {
    @Resource
    private ClientV4 clientV4;

    /**
     * 较稳定的随机数
     */
    public static final float STABLE_TEMPERATURE = 0.05f;
    /**
     * 较不稳定的随机数
     */
    public static final float UNSTABLE_TEMPERATURE = 0.99f;

    /**
     * 构造请求
     * @param messages
     * @param stream
     * @param temperature
     * @return
     */
    public String doRequest(List<ChatMessage> messages,Boolean stream, Float temperature){
        ChatMessage result = null;
        try {
            //构造请求
            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(Constants.ModelChatGLM4)
                    .stream(stream)
                    .invokeMethod(Constants.invokeMethod)
                    .temperature(temperature)
                    .messages(messages)
                    .build();
            ModelApiResponse invokeEmbeddingsApi = clientV4.invokeModelApi(chatCompletionRequest);
            result = invokeEmbeddingsApi.getData().getChoices().get(0).getMessage();
            return result.getContent().toString();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR,e.getMessage());
        }
    }

    /**
     * 通用请求，简化了消息传递
     * @param systemMessage
     * @param userMessage
     * @param stream
     * @param temperature
     * @return
     */
    public String doRequest(String systemMessage, String userMessage, Boolean stream, Float temperature){
        //构造请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doRequest(messages,stream,temperature);
    }

    /**
     * 同步请求
     * @param systemMessage
     * @param userMessage
     * @param temperature
     * @return
     */
    public String doSyncRequest(String systemMessage, String userMessage, Float temperature){
        //构造请求
        return doRequest(systemMessage,userMessage,Boolean.FALSE,temperature);
    }

    /**
     * 同步调用，答案稳定
     * @param systemMessage
     * @param userMessage
     * @return
     */
    public String doSyncStableRequest(String systemMessage, String userMessage){
        //构造请求

        return doSyncRequest(systemMessage,userMessage,STABLE_TEMPERATURE);
    }

    /**
     * 同步调用，答案不稳定
     */
    public String doSyncUnStableRequest(String systemMessage, String userMessage){
        //构造请求
        return doSyncRequest(systemMessage,userMessage,UNSTABLE_TEMPERATURE);
    }

    /**
     * 构造流式请求
     * @param messages
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doStreamRequest(List<ChatMessage> messages, Float temperature){
        ChatMessage result = null;
        try {
            //构造请求
            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                    .model(Constants.ModelChatGLM4)
                    .stream(Boolean.TRUE)
                    .invokeMethod(Constants.invokeMethod)
                    .temperature(temperature)
                    .messages(messages)
                    .build();
            ModelApiResponse invokeEmbeddingsApi = clientV4.invokeModelApi(chatCompletionRequest);
            return invokeEmbeddingsApi.getFlowable();
        } catch (Exception e) {
            e.printStackTrace();
            throw new BusinessException(ErrorCode.SYSTEM_ERROR,e.getMessage());
        }
    }

    /**
     * 通用流式请求，简化了消息传递
     * @param systemMessage
     * @param userMessage
     * @param temperature
     * @return
     */
    public Flowable<ModelData> doStreamRequest(String systemMessage, String userMessage, Float temperature){
        //构造请求
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        return doStreamRequest(messages,temperature);
    }
}
